package kimi

import (
	"context"
	"encoding/json"
	"errors"
	"log"
	"log/slog"
	"net/http"
	"quipus/internal/llm"
	kimi "quipus/internal/llm/provider/kimi/client"
	"time"

	mark3labs "github.com/mark3labs/mcp-go/mcp"
)

type moonshot struct {
	baseUrl string
	key     string
	client  *http.Client
	log     func(ctx context.Context, caller string, request *http.Request, response *http.Response, elapse time.Duration)
}

type KimiClient struct {
	client kimi.Client[*moonshot]
	model  string
}

type kimiAuthInfo struct {
	SecretKey string `json:"secret_key"`
}

var (
	Logger           = moonshot{}
	CustomHTTPClient = moonshot{}
)

func (m moonshot) BaseUrl() string      { return m.baseUrl }
func (m moonshot) Key() string          { return m.key }
func (m moonshot) Client() *http.Client { return m.client }

func (m moonshot) Log(ctx context.Context, caller string, request *http.Request, response *http.Response, elapse time.Duration) {
	m.log(ctx, caller, request, response, elapse)
}

func NewClient(config *llm.LLMClientConfig) (*KimiClient, error) {
	authInfo := kimiAuthInfo{}
	err := json.Unmarshal([]byte(config.AuthInfo), &authInfo)
	if err != nil {
		return nil, err
	}

	client := kimi.NewClient[*moonshot](&moonshot{
		baseUrl: config.BaseUrl,
		key:     authInfo.SecretKey,
		client:  http.DefaultClient,
		log: func(ctx context.Context, caller string, request *http.Request, response *http.Response, elapse time.Duration) {
			log.Printf("[%s] %s %s", caller, request.URL, elapse)
		},
	})
	c := &KimiClient{client: client, model: config.Model}
	return c, nil
}

func (c *KimiClient) GetProviderName() string {
	return "kimi"
}

func (c *KimiClient) Chat(ctx context.Context, req *llm.ChatReq) (err error) {
	if req == nil {
		return nil
	}

	model := req.Model
	msgs := req.Msgs

	if len(msgs) < 1 {
		return errors.New("message is null")
	}

	if len(model) < 1 {
		model = c.model
	}

	if len(model) < 1 {
		return errors.New("model is null")
	}

	messages := []*kimi.Message{}
	for _, m := range msgs {
		messages = append(messages, &kimi.Message{
			Role:    m.Role,
			Content: &kimi.Content{Text: m.Content},
		})
	}

	reqMsg := &kimi.ChatCompletionStreamRequest{
		Model:       model,
		Messages:    messages,
		MaxTokens:   4096,
		N:           1,
		Temperature: "0.3",
	}

	// completion, err := c.client.CreateChatCompletion(ctx, reqMsg)
	stream, err := c.client.CreateChatCompletionStream(ctx, reqMsg)

	if err != nil {
		return err
	}
loop:
	for {
		select {
		case chunk := <-stream.C:
			choices := chunk.Choices
			for _, choice := range choices {
				select {
				case req.RspChan <- choice.Delta.Content.Text:
					// slog.Info("get response: ", choice.Delta.Content)
				case <-ctx.Done():
					slog.Info("kimi recv quit signal")
					break loop
				}
			}
		default:
		}
	}
	req.CompleteChan <- true
	return nil
}

func (c *KimiClient) Complete(ctx context.Context, req *llm.CompleteReq) (llm.Message, error) {
	if req == nil {
		return nil, nil
	}

	model := req.Model
	msgs := req.Msgs

	if len(msgs) < 1 {
		return nil, errors.New("message is null")
	}

	if len(model) < 1 {
		model = c.model
	}

	if len(model) < 1 {
		return nil, errors.New("model is null")
	}

	messages := []*kimi.Message{}
	for _, m := range msgs {
		messages = append(messages, &kimi.Message{
			Role:    m.Role,
			Content: &kimi.Content{Text: m.Content},
		})
	}

	reqMsg := &kimi.ChatCompletionRequest{
		Model:       model,
		Messages:    messages,
		N:           1,
		Temperature: "0.3",
	}

	completion, err := c.client.CreateChatCompletion(ctx, reqMsg)
	if err != nil {
		return nil, err
	}

	if len(completion.Choices) < 1 {
		return nil, errors.New("no response")
	}

	txt := ""
	for _, choice := range completion.Choices {
		txt += choice.Message.Content.Text
	}
	return &KimiMessage{Resp: completion, Choice: &completion.Choices[0]}, nil
}

func convertTool(tools []mark3labs.Tool) []*kimi.Tool {
	var kTools []*kimi.Tool
	for _, tool := range tools {
		t := &kimi.Tool{
			Type:     kimi.ToolTypeFunction,
			Function: tool,
		}

		kTools = append(kTools, t)
	}

	return kTools
}
func convertSchema(schema mark3labs.ToolInputSchema) map[string]interface{} {
	// Ensure required is a valid array, defaulting to empty if nil
	required := schema.Required
	if required == nil {
		required = []string{}
	}

	return map[string]interface{}{
		"type":       schema.Type,
		"properties": schema.Properties,
		"required":   required,
	}
}

type KimiMessage struct {
	Resp   *kimi.Completion
	Choice *struct {
		Index        kimi.NullableType[int]    `json:"index"`
		Message      *kimi.Message             `json:"message"`
		FinishReason kimi.NullableType[string] `json:"finish_reason"`
		LogProbs     json.RawMessage           `json:"logprobs"`
	}
}

func (m *KimiMessage) GetRole() string {
	return string(m.Choice.Message.Role)
}

func (m *KimiMessage) GetContent() string {
	return m.Choice.Message.Content.Text
}

func (m *KimiMessage) GetToolCalls() []llm.ToolCall {
	var calls []llm.ToolCall
	for _, call := range m.Choice.Message.ToolCalls {
		calls = append(calls, &ToolCallWrapper{*call})
	}
	return calls
}

func (m *KimiMessage) IsToolResponse() bool {
	return m.Choice.Message.ToolCalls != nil && len(m.Choice.Message.ToolCalls) > 0
}

func (m *KimiMessage) GetToolResponseID() string {
	return m.Choice.Message.ToolCalls[0].ID
}

func (m *KimiMessage) GetUsage() (int, int) {
	return int(m.Resp.Usage.CompletionTokens), int(m.Resp.Usage.PromptTokens)
}

type ToolCallWrapper struct {
	Call kimi.ToolCall
}

func (t *ToolCallWrapper) GetID() string {
	return t.Call.ID
}

func (t *ToolCallWrapper) GetName() string {
	return t.Call.Function.Name
}

func (t *ToolCallWrapper) GetArguments() map[string]interface{} {
	var args map[string]interface{}
	if err := json.Unmarshal([]byte(t.Call.Function.Arguments), &args); err != nil {
		return make(map[string]interface{})
	}
	return args
}
